#!/usr/bin/python

from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC
import hmmtop,re
import ProjectBio
import pickle,os

if __name__=="__main__":
		
	tcdb = open("./tcdb")
	fastas = SeqIO.parse(tcdb,'fasta')
	fastas = ProjectBio.nr_dict(fastas)
	fastadb = {}
	def title_extract(string): #returns: (acc,tcid)
		string = string.split(" ")
		gi = string[0].split('|')[-1]
		tcid = string[1]
		return (gi,tcid)
	
	for fasta in fastas.values():
		(acc,tcid)=title_extract(fasta.description)
		fastadb.setdefault(tcid,[]).append(fasta.id)
	
	potential = dict([i for i in fastadb.items() if len(i[1])>1])
	# Create subset of FASTA files
	crazyids = []
	[crazyids.extend(i) for i in potential.values()]
	crazyids=list(set(crazyids))
	output=open("no_aux.faa",'wb')
	for crazy in crazyids:
		myseq = fastas[crazy]
		SeqIO.write(myseq,output,'fasta')
	output.close()
	# Check if HMMTOP has already been run?
	if os.path.exists("./no_aux.hmmtop") is True:
		tms = pickle.load(open("./no_aux.hmmtop"))
	else:
		ht = hmmtop.tools()
		ht.add_library("TCDB","./no_aux.faa")
		ht.scan_libraries()
		tms = ht.results['TCDB']
		pickle.dump(tms,open("./no_aux.hmmtop",'wb'))
	
	# Build TMS model!
	transporters = []
	for tcid,gis in potential.items():
		tmdict = {}
		for gi in gis:
			try:
				tmdict.setdefault(len(tms[gi]),[]).append(gi)
			except KeyError:
				continue
		keys = tmdict.keys()
		keys.sort(reverse=True)
		try:
			transporter = tmdict[keys[0]]
			transporters.append((tcid,transporter))
		except IndexError:
			# Build Length Model
			gis.sort(key=lambda x:len(fastas[x]),reverse=True)
			transporters.append((tcid,[gis[0]]))
	transporters = dict(transporters)
	#Build entire DB, selecting only transporters
	finaldb = []
	for tcid, gis in fastadb.items():
		if tcid in transporters.keys():
			finaldb.extend(transporters[tcid])
		else:
			finaldb.extend(gis)
	finaldb = list(set(finaldb))
	final_fasta = open("no_aux.faa",'wb')
	for gi in finaldb:
		entry = fastas[gi]
		entry.seq = Seq(re.sub('[^A-Za-z]','',str(entry.seq)),IUPAC.protein)
		SeqIO.write(entry,final_fasta,'fasta')
	
